package com.piece.core.db.context;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.piece.core.framework.constant.FrameWorkConstants;
import com.piece.core.framework.util.basic.SpringUtil;
import javax.sql.DataSource;

/**
 * 数据源切换处理
 */
public class DynamicDataSourceContextHolder {

    private static final ThreadLocal<String> CONTEXT_HOLDER = new TransmittableThreadLocal<>();

    /**
     * 设置数据源的变量
     */
    public static void setDateSourceType(String dsType) {
        CONTEXT_HOLDER.set(dsType);
    }

    /**
     * 获得数据源的变量
     */
    public static String getDateSourceType() {
        return CONTEXT_HOLDER.get();
    }

    /**
     * 清空数据源变量
     */
    public static void clearDateSourceType() {
        CONTEXT_HOLDER.remove();
    }

    public static DataSource getCurrentDataSource() {
        DataSource masterDataSource = (DataSource) SpringUtil.getBean("dataSource");
        DynamicDataSource dynamicDataSource = (DynamicDataSource) SpringUtil.getBean("dynamicDataSource");
        if (FrameWorkConstants.DATASOURCE_MASTER.equals(getDateSourceType())) {
            return masterDataSource;
        }
        return dynamicDataSource;
    }
}
